package com.yanghui.dtracing.core.opentracing.propagation;

import cn.hutool.core.util.StrUtil;
import com.yanghui.dtracing.core.opentracing.DTracerSpanContext;
import io.opentracing.SpanContext;
import io.opentracing.propagation.TextMapExtract;
import io.opentracing.propagation.TextMapInject;

import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

public class TextMapCodec implements Injector<TextMapInject>,Extractor<TextMapExtract> {

    @Override
    public void inject(SpanContext spanContext, TextMapInject carrier) {
        DTracerSpanContext context = (DTracerSpanContext)spanContext;
        carrier.put("traceId",context.toTraceId());
        carrier.put("spanId",context.toSpanId());
        carrier.put("sampled",String.valueOf(context.isSampled()));
        carrier.put("parentId",context.getParentId());
        Iterator<Map.Entry<String,String>> it = spanContext.baggageItems().iterator();
        while(it.hasNext()){
            Map.Entry<String,String> entry = it.next();
            carrier.put(entry.getKey(),entry.getValue());
        }
    }

    @Override
    public DTracerSpanContext extract(TextMapExtract carrier) {
        Iterator<Map.Entry<String,String>> iterator = carrier.iterator();
        String traceId = null;
        String spanId = null;
        Boolean sampled = true;
        String parentId = null;
        Map<String,String> baggages = new HashMap<>();
        while(iterator.hasNext()){
            Map.Entry<String,String> entry = iterator.next();
            if(entry.getKey().equals("traceId")){
                traceId = entry.getValue();
            }else if(entry.getKey().equals("spanId")){
                spanId = entry.getValue();
            }else if(entry.getKey().equals("sampled")){
                sampled = Boolean.valueOf(entry.getValue());
            }else if(entry.getKey().equals("parentId")){
                parentId = entry.getValue();
            }else {
                baggages.put(entry.getKey(),entry.getValue());
            }
        }
        if(StrUtil.isEmpty(traceId) || StrUtil.isEmpty(spanId)){
            return null;
        }
        DTracerSpanContext context = new DTracerSpanContext(traceId,spanId,parentId,sampled);
        context.addBaggageItems(baggages);
        return context;
    }
}
